#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
模板匹配应用程序 - 图像视图控件
实现图像显示、缩放、拖拽和选择区域功能
"""

import numpy as np
from PyQt5.QtWidgets import QWidget, QSizePolicy
from PyQt5.QtGui import QPainter, QPen, QBrush, QColor, QImage, QPixmap, QCursor
from PyQt5.QtCore import Qt, QRect, QPoint, QSize, pyqtSignal

class ImageView(QWidget):
    """
    图像显示控件
    支持图像显示、缩放、拖拽和选择矩形区域
    """
    
    def __init__(self, parent=None):
        """初始化图像视图控件"""
        super().__init__(parent)
        
        # 图像和显示相关属性
        self.image = None  # 原始图像
        self.pixmap = None  # 用于显示的图像
        self.zoom_factor = 1.0  # 缩放因子
        self.image_position = QPoint(0, 0)  # 图像位置
        
        # 拖拽相关属性
        self.is_dragging = False  # 是否正在拖拽
        self.last_drag_pos = QPoint()  # 上次拖拽位置
        
        # 选择矩形相关属性
        self.selection_rect = QRect()  # 选择矩形区域
        self.is_selecting = False  # 是否正在选择
        self.selection_start = QPoint()  # 选择起始点
        self.is_moving_selection = False  # 是否正在移动选择区域
        self.enable_selection = True  # 是否启用选择功能
        
        # 设置控件属性
        self.setMouseTracking(True)  # 开启鼠标跟踪
        self.setFocusPolicy(Qt.StrongFocus)  # 允许键盘焦点
        self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)  # 设置尺寸策略
        self.setMinimumSize(200, 200)  # 设置最小尺寸
        
    def set_image(self, image):
        """
        设置要显示的图像
        
        参数:
            image: QImage、QPixmap或NumPy数组
        """
        if image is None:
            return
            
        # 根据图像类型进行转换
        if isinstance(image, QPixmap):
            self.pixmap = image
            self.image = image.toImage()
        elif isinstance(image, QImage):
            self.image = image
            self.pixmap = QPixmap.fromImage(image)
        elif isinstance(image, np.ndarray):
            # 从NumPy数组转换为QImage
            height, width = image.shape[:2]
            if len(image.shape) == 3:
                if image.shape[2] == 3:
                    # RGB图像
                    qimg = QImage(image.data, width, height, width * 3, QImage.Format_RGB888)
                elif image.shape[2] == 4:
                    # RGBA图像
                    qimg = QImage(image.data, width, height, width * 4, QImage.Format_RGBA8888)
            else:
                # 灰度图像
                qimg = QImage(image.data, width, height, width, QImage.Format_Grayscale8)
                
            self.image = qimg
            self.pixmap = QPixmap.fromImage(qimg)
        else:
            return  # 不支持的类型
        
        # 重置显示属性
        self.zoom_factor = 1.0
        self.image_position = QPoint(0, 0)
        self.selection_rect = QRect()  # 清除选择区域
        self.center_image()  # 居中显示图像
        self.update()  # 触发重绘
        
    def center_image(self):
        """居中显示图像"""
        if self.pixmap is None:
            return
            
        # 计算图像的缩放尺寸
        scaled_size = self.calculate_scaled_size()
        
        # 计算居中位置
        x = (self.width() - scaled_size.width()) // 2
        y = (self.height() - scaled_size.height()) // 2
        
        self.image_position = QPoint(x, y)
        
    def calculate_scaled_size(self):
        """计算缩放后的图像尺寸"""
        if self.pixmap is None:
            return QSize(0, 0)
            
        return QSize(
            int(self.pixmap.width() * self.zoom_factor),
            int(self.pixmap.height() * self.zoom_factor)
        )
        
    def paintEvent(self, event):
        """重写绘制事件"""
        if self.pixmap is None:
            return
            
        painter = QPainter(self)
        painter.setRenderHint(QPainter.SmoothPixmapTransform, True)
        painter.setRenderHint(QPainter.Antialiasing, True)
        
        # 计算缩放后的图像尺寸
        scaled_size = self.calculate_scaled_size()
        
        # 绘制图像
        dest_rect = QRect(self.image_position, scaled_size)
        painter.drawPixmap(dest_rect, self.pixmap)
        
        # 绘制选择区域（如果存在）
        if not self.selection_rect.isEmpty():
            # 将图像坐标转换为视图坐标
            display_rect = self.convert_to_display_rect(self.selection_rect)
            
            # 绘制半透明背景
            painter.setPen(QPen(QColor(255, 0, 0), 2, Qt.SolidLine))
            painter.setBrush(QBrush(QColor(0, 0, 255, 50)))  # 半透明蓝色
            painter.drawRect(display_rect)
            
            # 绘制调整手柄
            self.draw_selection_handles(painter, display_rect)
            
    def draw_selection_handles(self, painter, rect):
        """绘制选择区域的调整手柄"""
        handle_size = 8
        painter.setPen(QPen(Qt.white, 1, Qt.SolidLine))
        painter.setBrush(QBrush(Qt.white))
        
        # 绘制四个角的手柄
        painter.drawRect(rect.left() - handle_size // 2, rect.top() - handle_size // 2, handle_size, handle_size)
        painter.drawRect(rect.right() - handle_size // 2, rect.top() - handle_size // 2, handle_size, handle_size)
        painter.drawRect(rect.left() - handle_size // 2, rect.bottom() - handle_size // 2, handle_size, handle_size)
        painter.drawRect(rect.right() - handle_size // 2, rect.bottom() - handle_size // 2, handle_size, handle_size)
        
    def convert_to_display_rect(self, image_rect):
        """将图像坐标转换为显示坐标"""
        if self.pixmap is None:
            return QRect()
            
        # 计算缩放比例
        scale_x = self.calculate_scaled_size().width() / self.pixmap.width()
        scale_y = self.calculate_scaled_size().height() / self.pixmap.height()
        
        # 计算显示坐标
        return QRect(
            self.image_position.x() + int(image_rect.x() * scale_x),
            self.image_position.y() + int(image_rect.y() * scale_y),
            int(image_rect.width() * scale_x),
            int(image_rect.height() * scale_y)
        )
        
    def convert_to_image_rect(self, display_rect):
        """将显示坐标转换为图像坐标"""
        if self.pixmap is None:
            return QRect()
            
        # 计算缩放比例
        scale_x = self.pixmap.width() / self.calculate_scaled_size().width()
        scale_y = self.pixmap.height() / self.calculate_scaled_size().height()
        
        # 计算图像坐标
        return QRect(
            int((display_rect.x() - self.image_position.x()) * scale_x),
            int((display_rect.y() - self.image_position.y()) * scale_y),
            int(display_rect.width() * scale_x),
            int(display_rect.height() * scale_y)
        )
        
    def convert_to_image_point(self, display_point):
        """将显示坐标点转换为图像坐标点"""
        if self.pixmap is None:
            return QPoint()
            
        # 计算缩放比例
        scale_x = self.pixmap.width() / self.calculate_scaled_size().width()
        scale_y = self.pixmap.height() / self.calculate_scaled_size().height()
        
        # 计算图像坐标点
        return QPoint(
            int((display_point.x() - self.image_position.x()) * scale_x),
            int((display_point.y() - self.image_position.y()) * scale_y)
        )
        
    def mousePressEvent(self, event):
        """鼠标按下事件"""
        if self.pixmap is None:
            return
            
        # 计算图像显示区域
        scaled_size = self.calculate_scaled_size()
        dest_rect = QRect(self.image_position, scaled_size)
        
        # 检查是否在图像区域内
        if not dest_rect.contains(event.pos()):
            return
            
        if self.enable_selection and event.button() == Qt.LeftButton:
            # 检查是否点击了选择区域
            if not self.selection_rect.isEmpty():
                display_rect = self.convert_to_display_rect(self.selection_rect)
                
                # 检查是否点击了调整手柄
                if self.is_on_selection_handle(event.pos(), display_rect):
                    self.is_selecting = True
                    self.selection_start = event.pos()
                    return
                    
                # 检查是否点击了选择区域内部
                if display_rect.contains(event.pos()):
                    self.is_moving_selection = True
                    self.last_drag_pos = event.pos()
                    self.setCursor(Qt.SizeAllCursor)
                    return
            
            # 开始新的选择
            self.is_selecting = True
            self.selection_start = event.pos()
            
            # 转换坐标为图像坐标
            image_point = self.convert_to_image_point(event.pos())
            self.selection_rect = QRect(image_point, QSize(0, 0))
        elif event.button() == Qt.LeftButton:
            # 开始拖拽图像
            self.is_dragging = True
            self.last_drag_pos = event.pos()
            self.setCursor(Qt.ClosedHandCursor)
        
    def mouseMoveEvent(self, event):
        """鼠标移动事件"""
        if self.pixmap is None:
            return
            
        # 计算图像显示区域
        scaled_size = self.calculate_scaled_size()
        dest_rect = QRect(self.image_position, scaled_size)
        
        if self.is_selecting and self.enable_selection:
            # 更新选择区域
            current_pos = event.pos()
            image_start = self.convert_to_image_point(self.selection_start)
            image_end = self.convert_to_image_point(current_pos)
            
            # 创建矩形区域
            x = min(image_start.x(), image_end.x())
            y = min(image_start.y(), image_end.y())
            width = abs(image_end.x() - image_start.x())
            height = abs(image_end.y() - image_start.y())
            
            # 限制在图像范围内
            x = max(0, min(x, self.pixmap.width() - 1))
            y = max(0, min(y, self.pixmap.height() - 1))
            width = min(width, self.pixmap.width() - x)
            height = min(height, self.pixmap.height() - y)
            
            self.selection_rect = QRect(x, y, width, height)
            self.update()
        elif self.is_moving_selection and self.enable_selection:
            # 移动选择区域
            delta_x = event.pos().x() - self.last_drag_pos.x()
            delta_y = event.pos().y() - self.last_drag_pos.y()
            
            # 获取当前显示矩形
            display_rect = self.convert_to_display_rect(self.selection_rect)
            
            # 移动显示矩形
            display_rect.translate(delta_x, delta_y)
            
            # 转换回图像坐标
            self.selection_rect = self.convert_to_image_rect(display_rect)
            
            # 限制在图像范围内
            if self.selection_rect.left() < 0:
                self.selection_rect.moveLeft(0)
            if self.selection_rect.top() < 0:
                self.selection_rect.moveTop(0)
            if self.selection_rect.right() >= self.pixmap.width():
                self.selection_rect.moveRight(self.pixmap.width() - 1)
            if self.selection_rect.bottom() >= self.pixmap.height():
                self.selection_rect.moveBottom(self.pixmap.height() - 1)
                
            self.last_drag_pos = event.pos()
            self.update()
        elif self.is_dragging:
            # 拖拽图像
            delta_x = event.pos().x() - self.last_drag_pos.x()
            delta_y = event.pos().y() - self.last_drag_pos.y()
            
            self.image_position.setX(self.image_position.x() + delta_x)
            self.image_position.setY(self.image_position.y() + delta_y)
            
            self.last_drag_pos = event.pos()
            self.update()
        elif self.enable_selection and not self.selection_rect.isEmpty():
            # 检查鼠标是否在选择区域上
            display_rect = self.convert_to_display_rect(self.selection_rect)
            
            if self.is_on_selection_handle(event.pos(), display_rect):
                self.setCursor(Qt.SizeFDiagCursor)
            elif display_rect.contains(event.pos()):
                self.setCursor(Qt.SizeAllCursor)
            else:
                self.setCursor(Qt.ArrowCursor)
        
    def mouseReleaseEvent(self, event):
        """鼠标释放事件"""
        self.is_dragging = False
        self.is_selecting = False
        self.is_moving_selection = False
        self.setCursor(Qt.ArrowCursor)
        
    def wheelEvent(self, event):
        """鼠标滚轮事件处理"""
        if self.pixmap is None:
            return
            
        # 计算缩放因子调整
        zoom_delta = 0.1 if event.angleDelta().y() > 0 else -0.1
        new_zoom = self.zoom_factor + zoom_delta
        
        # 限制缩放范围
        if new_zoom < 0.1:
            new_zoom = 0.1
        elif new_zoom > 10.0:
            new_zoom = 10.0
            
        # 以鼠标位置为中心缩放
        mouse_pos = event.pos()
        rel_x = (mouse_pos.x() - self.image_position.x()) / self.calculate_scaled_size().width()
        rel_y = (mouse_pos.y() - self.image_position.y()) / self.calculate_scaled_size().height()
        
        # 更新缩放因子
        self.zoom_factor = new_zoom
        
        # 计算新的图像位置
        new_scaled_size = self.calculate_scaled_size()
        new_x = mouse_pos.x() - rel_x * new_scaled_size.width()
        new_y = mouse_pos.y() - rel_y * new_scaled_size.height()
        
        self.image_position = QPoint(int(new_x), int(new_y))
        self.update()
        
    def mouseDoubleClickEvent(self, event):
        """鼠标双击事件处理"""
        # 重置缩放和位置
        self.zoom_factor = 1.0
        self.center_image()
        self.update()
        
    def is_on_selection_handle(self, point, rect):
        """检查点是否在选择区域的调整手柄上"""
        handle_size = 8
        
        # 检查四个角
        return (
            self.is_point_in_handle(point, rect.left(), rect.top(), handle_size) or
            self.is_point_in_handle(point, rect.right(), rect.top(), handle_size) or
            self.is_point_in_handle(point, rect.left(), rect.bottom(), handle_size) or
            self.is_point_in_handle(point, rect.right(), rect.bottom(), handle_size)
        )
        
    def is_point_in_handle(self, point, handle_x, handle_y, size):
        """检查点是否在手柄内"""
        return (
            point.x() >= handle_x - size // 2 and 
            point.x() <= handle_x + size // 2 and
            point.y() >= handle_y - size // 2 and 
            point.y() <= handle_y + size // 2
        )
        
    def crop_selected_region(self):
        """截取选中的区域"""
        if self.pixmap is None or self.selection_rect.isEmpty():
            return None
            
        # 从原始QPixmap裁剪出所选区域
        return self.pixmap.copy(self.selection_rect)
        
    def clear_selection(self):
        """清除选择区域"""
        self.selection_rect = QRect()
        self.update()


class TemplateView(ImageView):
    """
    模板图像视图控件
    继承自ImageView，但禁用选择功能，仅用于显示模板
    """
    
    def __init__(self, parent=None):
        """初始化模板视图控件"""
        super().__init__(parent)
        self.enable_selection = False  # 禁用选择功能
